'''
完成适配，待测试。
参考链接：
1. https://colab.research.google.com/drive/1N0Sf9yn8Tjs5gMJv-rez-0hzxBUDK3xK?usp=sharing#scrollTo=yl5Qoz2HzSUq
2. [Unsloth微调DeepSeek-R1蒸馏模型 - 构建医疗专家模型_哔哩哔哩_bilibili](https://www.bilibili.com/video/BV1qVNRenEBX/?spm_id_from=333.1387.favlist.content.click&vd_source=4d2a71ad349005831356ddd323e9568b)
3. [FreedomIntelligence/medical-o1-reasoning-SFT · Datasets at Hugging Face](https://huggingface.co/datasets/FreedomIntelligence/medical-o1-reasoning-SFT)

微调通过，但是在保存模型是失败，应该是有用到llama.cpp的代码导致的，错误信息如下：Unsloth: ##### The current model auto adds a BOS token.
Unsloth: ##### Your chat template has a BOS token. We shall remove it temporarily.
make: Entering directory '/home/lutery/projects/python/unsloth-total/llama.cpp'
Makefile:2: *** The Makefile build is deprecated. Use the CMake build instead. For more details, see https://github.com/ggml-org/llama.cpp/blob/master/docs/build.md.  Stop.
make: Leaving directory '/home/lutery/projects/python/unsloth-total/llama.cpp'
sh: 1: cmake: not found
Traceback (most recent call last):
  File "/home/lutery/projects/python/unsloth-total/deepseek_8b/deepseek_8b.py", line 197, in <module>
    main()
  File "/home/lutery/projects/python/unsloth-total/deepseek_8b/deepseek_8b.py", line 193, in main
    save_model_gguf_format(model, tokenizer)
  File "/home/lutery/projects/python/unsloth-total/deepseek_8b/deepseek_8b.py", line 64, in save_model_gguf_format
    if True: model.save_pretrained_gguf("model", tokenizer, quantization_method = "q4_k_m")
  File "/home/lutery/anaconda3/envs/deepseek/lib/python3.10/site-packages/unsloth/save.py", line 1814, in unsloth_save_pretrained_gguf
    makefile = install_llama_cpp_make_non_blocking()
  File "/home/lutery/anaconda3/envs/deepseek/lib/python3.10/site-packages/unsloth/save.py", line 786, in install_llama_cpp_make_non_blocking
    raise RuntimeError(f"*** Unsloth: Failed compiling llama.cpp using os.system(...) with error {check}. Please report this ASAP!")
RuntimeError: *** Unsloth: Failed compiling llama.cpp using os.system(...) with error 32512. Please report this ASAP!


本代码的训练目标：
本代码的训练目标是采用一种续写的方式实现问题的回答，
即让模型学习整体问题-推理-回答的完整文本，在完成训练后，输入给模型
的仅包含问题，让模型续写推理-回答，然后裁剪出推理和回答的部分
'''

from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported
import torch
from datasets import load_dataset
from huggingface_hub import create_repo

import os
# https://blog.csdn.net/zerolea/article/details/147224720
os.environ["http_proxy"] = "http://192.168.199.120:10802"
os.environ["https_proxy"] = "http://192.168.199.120:10802"

def test_model(model, tokenizer, prompt_style, question, device):
    FastLanguageModel.for_inference(model=model) # 将模型转换为推理模式，类似pytorch的eval模式，具体看md文件
    input = tokenizer([prompt_style.format(question, "")], return_tensors="pt").to(device)
    # 将输入的问题转换为适合模型输入的形式
    output = model.generate(
        input_ids=input.input_ids, # 提示词
        attention_mask=input.attention_mask, # 注意力掩码，主要是模型告诉1 代表是真实的词元，需要关注；0 代表是填充 (padding) 词元，应该忽略
        max_new_tokens=2000, # 设置模型输出的最大长度限制，防止无限循环
        do_sample=True, # 是否使用采样，True表示使用概率随机采样，False表示最大贪婪搜索，对于聊天、写作等创造性任务，强烈建议设置为 True。对于需要确定性答案的任务（如翻译、代码补全），可以考虑 False
        temperature=0.7, # 控制采样的随机性，值越高，输出越随机；值越低，输出越确定
        top_p=0.9, # 限制采样时考虑的词汇范围。只有在 do_sample=True 时才生效，它从概率最高的词开始，累加它们的概率，直到总和达到 top_p 的值。模型只会在这个累加出来的“核心”词汇表里进行采样
        use_cache=True, # 会将前面步骤的注意力计算结果（即键和值）缓存起来，下一步只计算当前新词的注意力结果，然后与缓存拼接即可，极大地加速了生成过程,Unsloth 的 FastLanguageModel.for_inference() 会自动优化这个缓存机制
    )
    response = tokenizer.batch_decode(output)
    print(response[0])


def get_datasets():
    # Login using e.g. `huggingface-cli login` to access this dataset
    # https://huggingface.co/datasets/FreedomIntelligence/medical-o1-reasoning-SFT
    # 看来可以直接通过huggingface上面的项目名字加载对应的数据集
    # 第二个参数表示加载的数据子集，可以在huggingface上面查看数据集的子集名称，不填应该是加载所有的数据
    ds = load_dataset("FreedomIntelligence/medical-o1-reasoning-SFT", "zh")
    print("Loaded dataset:", ds.column_names)
    return ds


def formatting_prompts_func(examples, prompt_style, EOS_TOKEN):
    '''
    将训练数据中的数据进行适当的转换，
    在本例中，提取其中的问题、推理、回答组合成一个字符串，
    并添加结束符号（EOS_TOKEN）以便模型理解文本的结束。
    添加到text标签中，并反馈，这样处理后就只有text列了
    '''
    inputs = examples['Question']
    cots = examples['Complex_CoT']
    outputs = examples['Response']
    texts = []
    for input, cot, output in zip(inputs, cots, outputs):
        text = prompt_style.format(input, cot, output) + EOS_TOKEN
        texts.append(text)
    return {"text": texts}


def save_model_gguf_format(model, tokenizer):
    # Save to 8bit Q8_0
    if False: model.save_pretrained_gguf("model", tokenizer,)

    # Save to 16bit GGUF
    if False: model.save_pretrained_gguf("model_f16", tokenizer, quantization_method = "f16")

    # Save to q4_k_m GGUF
    if True: model.save_pretrained_gguf("model", tokenizer, quantization_method = "q4_k_m")


def upload_huggingface_repo(model, tokenizer, repo_name="wyang14/medical-model", repo_id=None):
    create_repo(repo_name, token=HUGGINGFACE_TOKEN, exist_ok=True)
    model.push_to_hub_gguf(repo_name, tokenizer, token=HUGGINGFACE_TOKEN)

    # 上传到huggingfacehou可以使用ollama运行加载
    # ollama run hf.co/{username}/{repository}:{quantization}
    '''
    curl http://localhost:11434/api/chat -d '{ \
      "model": "", \
      "messages": [ \
        { "role": "user", "content": "一个患有急性阑尾炎的病人已经发病5天，腹痛稍有减轻但仍然发热，在体检时发现右下腹有压痛的包块，此时应如何处理？" } \
      }'
    '''


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"

    max_sequence_length = 2048
    dtype = None
    load_in_4bit = True

    # 可以在官网找到更多的模型
    # https://huggingface.co/unsloth/DeepSeek-R1-Distill-Qwen-1.5B-unsloth-bnb-4bit
    # https://docs.unsloth.ai/get-started/all-our-models#deepseek-models
    # 其中model是qwen的模型，tokenizer是qwen的分词器，均由模型自带
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name = "unsloth/DeepSeek-R1-Distill-Qwen-1.5B-unsloth-bnb-4bit",
        max_seq_length = max_sequence_length, # 定义模型的上下文长度，超过则阶段，不同的模型有不同的参数，另一个是max_sequence_length，如果错误则换一个试试
        dtype = dtype, # 设置加载模型计算时使用的浮点数类型，None表示自动选择合适的类型(比如硬件支持优先选择bfloat16，不支持则选择float16)，可以和权重文件本身的原始精度不匹配
                        # 比如使用4bit量化模型时，dtype可以设置为torch.float16或torch.bfloat16，这样的收益是最好的，也可以使用float32但是会占用更多的内存且收益较低
        load_in_4bit = load_in_4bit)
    # load_in_4bit表示会将模型精度及性能量化加载到4bit，如果为false
    # 这个时候加载的模型权重则以实际的模型精度存储，然后在转换到dtype的精度
    # 根据dtype的制定的及精度，模型的加载后的权重会自动转换到制定的dtype类型，越接近原始模型的精度，性能越好
    # 越远离的收益越差永远不要将一个以半精度（fp16/bf16）存储的模型加载为全精度（fp32），这只会浪费资源。


    prompt_style = """Below is an instruction that describes a task, paired with an input that provides further context.
    Write a response that appropriately completes the request.
    Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.

    ### Instruction:
    You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning.
    Please answer the following medical question.

    ### Question:
    {}

    ### Response:
    <think>{}"""


    train_prompt_style = """
    Below is an instruction that describes a task, paired with an input that provides further context.
    Write a response that appropriately completes the request.

    Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.


    ### Instruction:
    You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning.
    Please answer the following medical question.

    ### Question:
    {}

    ### Response:
    <think>
    {}
    </think>
    {}"""

    question = "一个患有急性阑尾炎的病人已经发病5天，腹痛稍有减轻但仍然发热，在体检时发现右下腹有压痛的包块，此时应如何处理？"

    # test_model(model, tokenizer, prompt_style, question, device)
    # 我无法直接访问或查看医院的具体信息或 patient data。如果需要帮助，建议联系医院或相关医疗团队，他们有处理此类情 况的专业知识和技能。<｜end▁of▁sentence｜>

    EOS_TOKEN = tokenizer.eos_token if tokenizer.eos_token else "<|endoftext|>" # 获取分词器的结束符号，如果没有则使用默认的"<|endoftext|>"，这个符号在模型训练时会被用来标记文本的结束
    dataset = get_datasets()
    dataset = dataset.map(formatting_prompts_func, batched=True, fn_kwargs={"prompt_style": train_prompt_style, "EOS_TOKEN": EOS_TOKEN})
    print(f"Dataset first text is {dataset['train']['text'][0]}")

    # 返回一个包含了这些可训练 LoRA 适配器的新模型对象。在后续的训练中，只有这些新增的、百万级别的 LoRA 参数会被更新，而原始的数十亿参数保持不变
    # 
    model = FastLanguageModel.get_peft_model(
        model=model,
        r = 16, # LoRA 的秩 (Rank)。这是 LoRA 中最重要的超参数，越大，可训练参数越多，模型的拟合能力越强；反之，越小，可训练参数越少，模型的拟合能力越弱。通常建议从 8 或 16 开始尝试
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        # 指定要将 LoRA 适配器注入到模型的哪些层中，您提供的列表 ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] 是一个非常全面的配置，它覆盖了注意力机制中的所有四个线性层 (q, k, v, o) 
        lora_alpha=16, # LoRA 的缩放因子 (Alpha)，权衡原始权重和训练权重之间的比例，通常做法设置为r或者r的两倍
        lora_dropout=0, # 在 LoRA 适配器的权重上应用 Dropout，作为一种正则化手段，防止过拟合，0表示不是哦给你dropout
        bias="none",# 制偏置项 (bias) 是否可训练,none表示不训练偏置项，lora_only表示训练lora偏置项,all表示训练所有偏置项
        use_gradient_checkpointing='unsloth', # 这是一种用计算时间换取显存的强大技术。它在反向传播时不会保存所有的中间激活值，而是在需要时重新计算它们。这可以大幅降低显存占用，使得你可以用更大的批次大小或更长的序列长度进行训练
        random_state=3407, # 控制 LoRA 权重的初始化。设置一个固定的种子可以确保每次运行代码时，LoRA 矩阵的初始值都完全相同，从而使得训练过程可以完全复现
        use_rslora=False, # 是否使用 Rank-Stabilized LoRA。这是一种改进的 LoRA，通过调整 lora_alpha 和 r 的关系来稳定训练。Unsloth 的实现中，它会自动将 lora_alpha 设置为 r 的平方根
        loftq_config=None # 是否使用 LoftQ，一种在量化模型上进行 LoRA 初始化的技术，可以进一步提升性能
        # 些是更高级的 PEFT 技术。对于初次尝试，将它们保持为默认值（False 和 None）是完全可以的
    )
    # 对于target_modules参数可以通过一下几种方式找到实际的参数
    # 1. print(model)输出模型层的名字
    # 2. 查阅文档找到模型层相关的说明知道层的名字

    trainer = SFTTrainer(
        model=model, # 通过 FastLanguageModel.get_peft_model() 获取的模型，其中增加了可以微调的参数部分
        tokenizer=tokenizer, # 分词器，用于将文本转换为模型可以理解的输入格式
        train_dataset=dataset["train"], # 训练数据集，通常是一个包含文本和标签的 Dataset 对象
        dataset_text_field="text", # 指定数据集中哪一列包含训练文本
        max_seq_length=max_sequence_length,
        # 作用: 定义模型处理的最大序列长度，通常需要与模型加载时设置的一致。
        # 注意: 如果训练样本超过此长度，会被截断；低于此长度则会被填充。
        dataset_num_proc=2, # 数据集处理的并行进程数，通常设置为 CPU 核心数的 1/2 或 1/4，最多不应超过 CPU 核心数
        packing=False, # 控制是否将多个短样本"打包"到一个序列中
        # True: 多个短样本会被连接到一起，形成单个训练序列，提高训练效率。
        # False(您的选择): 每个样本单独处理，不合并。这更适合需要严格保持样本边界的情况。
        # 使用场景: 当您的训练数据由结构化的对话或需要严格按格式学习时，设为 False 更好；如果是简短的文本片段，设为 True 可以提高训练效率。
        args=TrainingArguments(
            per_device_train_batch_size=2, # 每个 GPU 设备一次处理的样本数
            gradient_accumulation_steps=4, # 在执行优化器更新前，累积多少个批次的梯度
            warmup_steps=5, # 学习率逐渐增加的步数，有助于稳定早期训练,通常设为总训练步数的 5-10%
            max_steps=60, # 训练的总步数
            learning_rate=2e-4, # 优化器的学习率
            fp16=not is_bfloat16_supported(), # 控制训练使用的浮点精度
            bf16=is_bfloat16_supported(), # 控制训练使用的浮点精度，根据是否支持 bfloat16 来设置
            logging_steps=1, # 每隔多少步记录一次日志
            optim="adamw_8bit", # 选择优化器，adamw_8bit: 比传统 AdamW 优化器更节省显存的 8 位实现，由 bitsandbytes 库提供
            weight_decay=0.01, # 权重衰减系数，用于防止过拟合
            lr_scheduler_type="linear", # 学习率调度器类型，linear: 学习率线性衰减,: 从初始学习率线性降低到 0，是微调的常用选择
            seed=3047, # 随机种子，用于确保实验的可重复性
            output_dir="outputs", # 保存检查点和训练结果的目录
            report_to="none", # 指定要向哪些平台报告训练指标， "none" 表示不报告；其他选择包括 "tensorboard", "wandb" 等
        )
    )

    trainer_stats = trainer.train()
    print("Training stats:", trainer_stats)

    test_model(
        model=model,
        tokenizer=tokenizer,
        prompt_style=prompt_style,
        question=question,
        device=device
    )   
    '''
    微调后
     ### Instruction:
    You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning.
    Please answer the following medical question.

    ### Question:
    一个患有急性阑尾炎的病人已经发病5天，腹痛稍有减轻但仍然发热，在体检时发现右下腹有压痛的包块，此时应如何处理？

    ### Response:
    <think>
    看来这位病人确实有点不好的症状，急性阑尾炎已经持续了5天，腹痛又稍微减轻了一点，但还有发热。这让我想到，他应该是在消化道里有炎症，所以要看看消化道有没有什么问题。根据这个情况，我得先看看他的胃部是不是有压力，因为胃里有压力，通常会导致胃痛。

然后，我再想到胃里有压力，可能会导致胃炎，所以检查胃镜应该是个不错的主意。如果发现胃里有压力，胃镜可以帮我们找到问题所在。如果发现胃里有压力，可能需要做一些胃部手术，比如胃切除术。

不过，如果胃里没有压力，那就得看其他原因，比如消化道的炎症，像胆囊炎或者胆囊炎。在胆囊炎的情况下，我需要先做胆管镜，看看胆囊是不是有什么问题。如果胆囊有囊膜炎，就得进行胆管手术，比如胆囊切除术。

然后，再想，如果这两种情况都排除了，可能还要考虑其他原因，比如胆囊炎或者胆囊炎，因为胆囊炎是胆囊炎的一种。如果胆囊没有压力，可能需要考虑胆管炎，然后才需要做胆管镜和手术。

总之，这个病人得先检查胃里有没有压力，如果有的话，先做胃镜，看看胃里有无压力。如果没压力，再做胆管镜，看看胆囊有没有囊膜炎，如果有，再手术。如果胆囊没有压力，但有胆囊炎，可能需要更深入的检查。

啊，看来我得再理清楚这些步骤，确保每一步都合理，不会遗漏什么。毕竟，医生要根据病人的情况来决定接下来的检查和手术。
    </think>
    根据病人的情况，首先应该检查胃里是否有压力。如果发现胃里有压力，应该进行胃镜检查，以确定胃里有无压力。如果胃里没有压力，需要进一步考虑其他原因，如胆囊炎。在这种情况下，应该先检查胆囊是否有囊膜炎。如果胆囊有囊膜炎，可能需要进行胆管镜检查，以确定胆囊的状况。如果胆囊没有压力，但有胆囊炎，可能需要更深入的检查，如胆管镜，以确定胆囊的状况。如果胆囊没有压力，但有胆囊炎，可能需要考虑其他原因，如胆管炎，然后进行相应的检查和手术。<｜end▁of▁sentence｜>
    '''
    save_model_gguf_format(model, tokenizer)


if __name__ == "__main__":
    main()